#include <vitasdk.h>
#include <taihen.h>

#include "scr_print.h"
#include "menu.h"
#include "config.h"
#include "utils.h"

typedef struct CtrlEnttry
{
    int key;
    int32_t *value;
    uint64_t last_micros;
} CtrlEnttry;

#define MAX_HOOKS_NUM 13
static SceUID s_hooks[MAX_HOOKS_NUM];
static tai_hook_ref_t s_hook_refs[MAX_HOOKS_NUM];

static SceKernelLwMutexWork s_draw_mutex;
char g_titleid[32]; 

static SceButtons ctrl_buttons, ctrlex2_buttons;

static CtrlEnttry ctrl_entries[] = {
    {SCE_CTRL_LEFT, &setting_config.button_left, 0},
    {SCE_CTRL_UP, &setting_config.button_up, 0},
    {SCE_CTRL_RIGHT, &setting_config.button_right, 0},
    {SCE_CTRL_DOWN, &setting_config.button_down, 0},
    {SCE_CTRL_CROSS, &setting_config.button_cross, 0},
    {SCE_CTRL_CIRCLE, &setting_config.button_circle, 0},
    {SCE_CTRL_SQUARE, &setting_config.button_square, 0},
    {SCE_CTRL_TRIANGLE, &setting_config.button_triangle, 0},
    {SCE_CTRL_SELECT, &setting_config.button_select, 0},
    {SCE_CTRL_START, &setting_config.button_start, 0},
    {SCE_CTRL_LTRIGGER, &setting_config.button_l, 0},
    {SCE_CTRL_RTRIGGER, &setting_config.button_r, 0},
};

#define N_CTRL_ENTRIES (sizeof(ctrl_entries) / sizeof(CtrlEnttry))

static CtrlEnttry ctrlex2_entries[] = {
    {SCE_CTRL_LEFT, &setting_config.button_left, 0},
    {SCE_CTRL_UP, &setting_config.button_up, 0},
    {SCE_CTRL_RIGHT, &setting_config.button_right, 0},
    {SCE_CTRL_DOWN, &setting_config.button_down, 0},
    {SCE_CTRL_CROSS, &setting_config.button_cross, 0},
    {SCE_CTRL_CIRCLE, &setting_config.button_circle, 0},
    {SCE_CTRL_SQUARE, &setting_config.button_square, 0},
    {SCE_CTRL_TRIANGLE, &setting_config.button_triangle, 0},
    {SCE_CTRL_SELECT, &setting_config.button_select, 0},
    {SCE_CTRL_START, &setting_config.button_start, 0},
    {SCE_CTRL_L1, &setting_config.button_l, 0},
    {SCE_CTRL_R1, &setting_config.button_r, 0},
    {SCE_CTRL_L2, &setting_config.button_l2, 0},
    {SCE_CTRL_R2, &setting_config.button_r2, 0},
    {SCE_CTRL_L3, &setting_config.button_l3, 0},
    {SCE_CTRL_R3, &setting_config.button_r3, 0},
};

#define N_CTRLEX2_ENTRIES (sizeof(ctrlex2_entries) / sizeof(CtrlEnttry))

static void cleanKey(uint32_t *buttons, int key, int negative)
{
    key = negative ? ~key : key;
    *buttons &= ~key;
}

static void applyTurbo(SceCtrlData *ctrl, int negative, int is_ex2)
{
    uint64_t cur_micros = sceKernelGetProcessTimeWide();
    uint64_t interval_micros, delay_micros;
    SceButtons *buttons;
    CtrlEnttry *entries;
    int n_entries;

    if (is_ex2)
    {
        buttons = &ctrlex2_buttons;
        entries = ctrlex2_entries;
        n_entries = N_CTRLEX2_ENTRIES;
    }
    else
    {
        buttons = &ctrl_buttons;
        entries = ctrl_entries;
        n_entries = N_CTRL_ENTRIES;
    }

    readPadEx(buttons, ctrl, negative);

    int i;
    for (i = 0; i < n_entries; i++)
    {
        uint64_t value = *(entries[i].value);
        if (value > 0)
        {
            if (isKeyCurrentEx(buttons, entries[i].key))
            {
                if (!isKeyOldEx(buttons, entries[i].key))
                {
                    entries[i].last_micros = cur_micros;
                }
                else
                {
                    interval_micros = cur_micros - entries[i].last_micros;
                    delay_micros = (turbo_options.min + value) * 1000;
                    if (interval_micros < delay_micros)
                        cleanKey(&ctrl->buttons, entries[i].key, negative);
                    else
                        entries[i].last_micros = cur_micros;
                }
            }
        }
    }
}

static int sceDisplaySetFrameBuf_patched(const SceDisplayFrameBuf *pParam, int sync)
{
    if (isShowingMenu())
    {
        sceKernelLockLwMutex(&s_draw_mutex, 1, NULL);
        scrPrintSetFrameBuf(pParam);
        drawMenu();
        sceKernelUnlockLwMutex(&s_draw_mutex, 1);
    }
    return TAI_CONTINUE(int, s_hook_refs[0], pParam, sync);
}

static int ctrlPatched(tai_hook_ref_t hook_ref, int port, SceCtrlData *ctrl, int count, int negative, int is_ex2)
{
    if (hook_ref == 0)
        return 1;

    int ret = TAI_CONTINUE(int, hook_ref, port, ctrl, count);
    if (ret < 0 || (port != 0 && port != 1))
        return ret;

    if (!isShowingMenu())
    {
        readPad(ctrl, negative);
        if (isKeyCurrent(SCE_CTRL_START) && (isKeyCurrent(SCE_CTRL_LTRIGGER) || isKeyCurrent(SCE_CTRL_L1)))
            openMenu();
        else if (setting_config.enable_turbo)
            applyTurbo(ctrl, negative, is_ex2);
    }
    else
    {
        sceKernelLockLwMutex(&s_draw_mutex, 1, NULL);
        ctrlMenu(ctrl, negative);
        sceKernelUnlockLwMutex(&s_draw_mutex, 1);
    }

    return ret;
}

static int sceCtrlPeekBufferPositive_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[1], port, ctrl, count, 0, 0);
}

static int sceCtrlPeekBufferPositive2_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[2], port, ctrl, count, 0, 0);
}

static int sceCtrlReadBufferPositive_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[3], port, ctrl, count, 0, 0);
}

static int sceCtrlReadBufferPositive2_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[4], port, ctrl, count, 0, 0);
}

static int sceCtrlPeekBufferPositiveExt_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[5], port, ctrl, count, 0, 0);
}

static int sceCtrlPeekBufferPositiveExt2_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[6], port, ctrl, count, 0, 1);
}

static int sceCtrlReadBufferPositiveExt_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[7], port, ctrl, count, 0, 0);
}

static int sceCtrlReadBufferPositiveExt2_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[8], port, ctrl, count, 0, 1);
}

static int sceCtrlPeekBufferNegative_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[9], port, ctrl, count, 1, 0);
}

static int sceCtrlPeekBufferNegative2_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[10], port, ctrl, count, 1, 0);
}

static int sceCtrlReadBufferNegative_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[11], port, ctrl, count, 1, 0);
}

static int sceCtrlReadBufferNegative2_patched(int port, SceCtrlData *ctrl, int count)
{
    return ctrlPatched(s_hook_refs[12], port, ctrl, count, 1, 0);
}

void _start() __attribute__((weak, alias("module_start")));
int module_start(SceSize args, void *argp)
{
    sceKernelCreateLwMutex(&s_draw_mutex, "user_draw_mutex", 2, 0, NULL);
    sceAppMgrAppParamGetString(0, 12, g_titleid, 256);
    loadConfig();

    s_hooks[0] = taiHookFunctionImport(&s_hook_refs[0], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0x7A410B64, sceDisplaySetFrameBuf_patched);

    s_hooks[1] = taiHookFunctionImport(&s_hook_refs[1], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0xA9C3CED6, sceCtrlPeekBufferPositive_patched);
    s_hooks[2] = taiHookFunctionImport(&s_hook_refs[2], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0x15F81E8C, sceCtrlPeekBufferPositive2_patched);
    s_hooks[3] = taiHookFunctionImport(&s_hook_refs[3], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0x67E7AB83, sceCtrlReadBufferPositive_patched);
    s_hooks[4] = taiHookFunctionImport(&s_hook_refs[4], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0xC4226A3E, sceCtrlReadBufferPositive2_patched);
    s_hooks[5] = taiHookFunctionImport(&s_hook_refs[5], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0xA59454D3, sceCtrlPeekBufferPositiveExt_patched);
    s_hooks[6] = taiHookFunctionImport(&s_hook_refs[6], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0x860BF292, sceCtrlPeekBufferPositiveExt2_patched);
    s_hooks[7] = taiHookFunctionImport(&s_hook_refs[7], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0xE2D99296, sceCtrlReadBufferPositiveExt_patched);
    s_hooks[8] = taiHookFunctionImport(&s_hook_refs[8], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0xA7178860, sceCtrlReadBufferPositiveExt2_patched);
    s_hooks[9] = taiHookFunctionImport(&s_hook_refs[9], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0x104ED1A7, sceCtrlPeekBufferNegative_patched);
    s_hooks[10] = taiHookFunctionImport(&s_hook_refs[10], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0x81A89660, sceCtrlPeekBufferNegative2_patched);
    s_hooks[11] = taiHookFunctionImport(&s_hook_refs[11], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0x15F96FB0, sceCtrlReadBufferNegative_patched);
    s_hooks[12] = taiHookFunctionImport(&s_hook_refs[12], TAI_MAIN_MODULE, TAI_ANY_LIBRARY, 0x27A0C5FB, sceCtrlReadBufferNegative2_patched);

    return SCE_KERNEL_START_SUCCESS;
}

int module_stop(SceSize argc, const void *args)
{
    int i;
    for (i = 0; i < MAX_HOOKS_NUM; i++)
    {
        if (s_hooks[i] >= 0)
            taiHookRelease(s_hooks[i], s_hook_refs[i]);
    }

    sceKernelDeleteLwMutex(&s_draw_mutex);

    return SCE_KERNEL_STOP_SUCCESS;
}